﻿/*
	本文件功能：
	枚举模块代码
*/
#pragma once

#include "stdafx.h"
#include "utf8.h"

int EnmuDisk(OUT LIST_ENTRY* pDiskListHead){
	DbgPrint("Enmuing Disks...");
	LIST_ENTRY list_entry;
	PHANDLE CurrentDiskHandle;
	wchar_t DiskIDName[33];
	UNICODE_STRING DiskName;
	OBJECT_ATTRIBUTES DiskAttributes;
	IO_STATUS_BLOCK stat_blk;
	NTSTATUS status;
	DiskListEntry *elem;
	void *buffer = NULL;
	UINT DiskNum;
	for (UINT i =0;;i++){
		swprintf(DiskIDName,L"\\.\\PhysicalDriver%u",i);
		RtlInitUnicodeString(&DiskName,DiskIDName);
		InitializeObjectAttributes(&DiskAttributes,&DiskName,OBJ_CASE_INSENSITIVE|OBJ_KERNEL_HANDLE,NULL,NULL);
		status = ZwCreateFile(CurrentDiskHandle,FILE_GENERIC_READ | FILE_GENERIC_WRITE,&DiskAttributes,&stat_blk,NULL,FILE_ATTRIBUTE_NORMAL,FILE_SHARE_READ | FILE_SHARE_WRITE,FILE_OPEN_IF,FILE_NON_DIRECTORY_FILE |FILE_RANDOM_ACCESS |FILE_SYNCHRONOUS_IO_NONALERT,NULL,0);
		if (NT_SUCCESS(status)){
			elem = (DiskListEntry *)ExAllocatePoolWithTag(PagedPool,sizeof(DiskListEntry),'disk');
			elem->CurrentDiskHandle = CurrentDiskHandle;
			wcscpy(elem->DiskIDName,DiskIDName);
			while (!buffer) ExAllocatePoolWithTag(PagedPool,33*sizeof(wchar_t),'str');
			RtlInitEmptyUnicodeString(&(elem->DiskName),(PWCHAR)buffer,33);
			RtlCopyUnicodeString(&(elem->DiskName),&DiskName);
			InsertTailList(pDiskListHead,&elem->list_item);

		}else{
			DiskNum = i + 1;
			break;
		}
	}
	return DiskNum;
}

int EnmuUD(IN LIST_ENTRY *pDiskListHead,OUT LIST_ENTRY *pUDListHead,IN PDRIVER_OBJECT DriverObject){
	PAGED_CODE();
	DiskListEntry *elem;
	char flag[4];
	LARGE_INTEGER offest;
	offest.QuadPart =0x1b4;
	IO_STATUS_BLOCK io_stk;
	int disks = 0;
	UNICODE_STRING device;
	wchar_t buf[2][33];
	char letter = 'Z';
	device.Buffer =buf[0];
	device.Length = 0;
	UNICODE_STRING symlink;
	for (PLIST_ENTRY p = pDiskListHead->Flink;p != pDiskListHead;p = p->Flink){
		elem = CONTAINING_RECORD(p,DiskListEntry,list_item);
		if (!NT_SUCCESS(ZwReadFile(elem->CurrentDiskHandle,NULL,NULL,NULL,&io_stk,flag,4,&offest,NULL))){
			return -1;
		}
		if (RtlEqualMemory(flag,"FBBF",4))
		{
			elem->Volume = (FBFSVolume *)ExAllocatePool(NonPagedPool,sizeof(FBFSVolume));
			EnmuFileList(elem);
			KeInitializeMutex(&elem->Volume->FileListLock,0);
			/*swprintf(buf,"\\DosDevices\\%c:",letter);*/
			RtlStringCbPrintfW(buf[0],33,L"\\Device\\UDVolume%d",disks);
			RtlInitUnicodeString(&device,buf[0]);
			elem->Volume->VolumeDevice = (PDEVICE_OBJECT)ExAllocatePool(NonPagedPool,sizeof(DEVICE_OBJECT));
			IoCreateDevice(DriverObject,sizeof(UD_Ext_Flag),&device,FILE_DEVICE_DISK,FILE_READ_ONLY_VOLUME|FILE_REMOVABLE_MEDIA,FALSE,&elem->Volume->VolumeDevice);
			UD_Ext_Flag *Flag = (UD_Ext_Flag*)elem->Volume->VolumeDevice->DeviceExtension;
			Flag->DriverObject = DriverObject;
			Flag->Signature = Volume_Sign;
			RtlStringCbPrintfW(buf[1],33,L"\\DosDevices\\%c:",letter);
			RtlInitUnicodeString(&symlink,buf[1]);
			IoCreateUnprotectedSymbolicLink(&symlink,&device);
;			InsertTailList(pUDListHead,&elem->udlist_item);
			++disks;
			--letter;
			continue;
		}else{
			continue;
		}
	}
	return disks;
}

bool IfIsUD(wchar_t *DiskIDName){
	ASSERT(DiskIDName);
	IO_STATUS_BLOCK stat_blk;
	NTSTATUS status;
	char *buf = (char*)ExAllocatePool(NonPagedPool,5);
	UNICODE_STRING DiskName;
	DiskName.Buffer = DiskIDName;
	DiskName.Length = wcslen(DiskIDName);
	DiskName.MaximumLength = wcslen(DiskIDName)+1;
	char Magic[5] = "FBBF";
	LARGE_INTEGER offest;
	offest.QuadPart = 0x1b4;
	OBJECT_ATTRIBUTES obj;
	HANDLE handle;
	InitializeObjectAttributes(&obj,&DiskName,OBJ_KERNEL_HANDLE|OBJ_CASE_INSENSITIVE,NULL,NULL);
	if (!NT_SUCCESS(ZwCreateFile(&handle,FILE_GENERIC_READ | FILE_GENERIC_WRITE,&obj,&stat_blk,NULL,FILE_ATTRIBUTE_NORMAL,FILE_SHARE_READ | FILE_SHARE_WRITE,FILE_OPEN_IF,FILE_NON_DIRECTORY_FILE |FILE_RANDOM_ACCESS |FILE_SYNCHRONOUS_IO_NONALERT,NULL,0))){
		ExFreePool((PVOID)buf);
		ZwClose(handle);
		return false;
	}
	if (!NT_SUCCESS(ZwReadFile(handle,NULL,NULL,NULL,&stat_blk,(PVOID)buf,4,&offest,0))){
		ExFreePool((PVOID)buf);
		ZwClose(handle);
		return false;
	}
	buf[4] = '\0';
	if (!strcmp(buf,Magic))
	{
		ExFreePool((PVOID)buf);
		ZwClose(handle);
		return true;
	}
	ExFreePool((PVOID)buf);
	ZwClose(handle);
	return false;
}


struct BufStruct{
	PVOID buf;
	ULONG Length;
	LIST_ENTRY list_entry;
};


//void GetUDFileLists(INOUT PLIST_ENTRY udDeviceList){
//	LIST_ENTRY buf_list;
//	PAGED_CODE();
//	IO_STATUS_BLOCK io_stk;
//	NTSTATUS status;
//	char buffer[2];
//	BufStruct *buf;
//	PVOID temps[5];
//	FileRecord *record;
//	hextonum c;
//	LARGE_INTEGER offest;
//	ULONG List_length;
//	DiskListEntry *elem;
//	int i = 0;
//	unsigned int k =1;
//	UINT l = 0;
//	UINT m =0;
//	UINT filelistoffest;
//	UINT disknum = 0;
//	for (LIST_ENTRY *p = udDeviceList->Flink;p != udDeviceList;p = p->Flink,++disknum){
//		offest.QuadPart = 64*512;
//		elem = CONTAINING_RECORD(p,DiskListEntry,udlist_item);
//		InitializeListHead(&elem->Volume->FileListHead);
//		KeInitializeMutex(&elem->Volume->FileListLock,0);
//		status = ZwReadFile(elem->CurrentDiskHandle,NULL,NULL,NULL,&io_stk,(PVOID)buffer,2,&offest,NULL);
//		RtlCopyMemory(c.hexarray,buffer,8);
//		for (UINT j = 0;j<c.shortnum;i++){
//			buf = (BufStruct *)ExAllocatePoolWithTag(PagedPool,sizeof(BufStruct),'file');
//			buf->buf = ExAllocatePoolWithTag(PagedPool,510,'lstb');
//			buf->Length = 510;
//			InsertTailList(&buf_list,&buf->list_entry);
//		}
//		//读取记录
//		offest.QuadPart = (64*512)+(c.shortnum)*512;
//		for (PLIST_ENTRY p = buf_list.Flink;p != &buf_list;p = p->Flink){
//			BufStruct *hc = CONTAINING_RECORD(p,BufStruct,list_entry);
//			status = ZwReadFile(elem->CurrentDiskHandle,NULL,NULL,NULL,&io_stk,hc->buf,510,&offest,NULL);
//			offest.QuadPart += (510+2);
//		}
//		//分解记录
//		temps[1] = ExAllocatePool(PagedPool,(c.shortnum)*512);
//		BufStruct *pBuf;
//		//首先合并记录
//		for (PLIST_ENTRY p = buf_list.Flink;p != &buf_list;p = p->Flink){
//			pBuf = CONTAINING_RECORD(p,BufStruct,list_entry);
//			memcat(temps[1],pBuf->buf,l,pBuf->Length);
//			ExFreePool(pBuf->buf);
//			ExFreePool(pBuf);
//
//			l += 510;
//		}
//		temps[4] = temps[1];
//		char filelistLength;
//		//正式分解记录
//		for (;k<l;k+=510)
//		{
//			filelistLength = *((char*)(temps[1]));
//			temps[3] = ExAllocatePool(PagedPool,filelistLength);
//			memcpy(temps[3],temps[1],filelistLength);
//			record = GetARecord((char*)temps[3],(FileRecord*)ExAllocatePool(NonPagedPool,sizeof(FileRecord)));
//			ExFreePool(temps[3]);
//			temps[1] += (filelistLength + 2);
//			wchar_t devicename[256];
//			swprintf(devicename,L"\\Device\\HarddiskVolume%u",m);
//			//RtlInitUnicodeString(&record->OnDevice->,devicename);
//			InsertTailList(&elem->Volume->FileListHead,&record->list_item);
//		}
//		m++;
//		//最后，释放内存
//		ExFreePool(temps[4]);
//	}
//	return;
//}
//


void *memcat(void* Dst,const void* src,unsigned int offest,size_t srcLength){
	unsigned int start = (unsigned int)Dst;
	if ((src==NULL)||(Dst==NULL)||(srcLength==0))
	{
		return NULL;
	}
	
	start += offest;
	void* st= (void*)start;
	RtlCopyMemory(st,src,srcLength);
	return st;
}

void EnmuFileList(DiskListEntry *Disk){
	__try{
		LIST_ENTRY buf_list;
		InitializeListHead(&buf_list);
		InitializeListHead(&Disk->Volume->FileListHead);
		PAGED_CODE();
		IO_STATUS_BLOCK io_stk;
		NTSTATUS status;
		char buffer[2];
		BufStruct *buf;
		PVOID temps[5];
		hextonum a;
		LARGE_INTEGER offest;
		ULONG List_length;
		int i = 0;
		unsigned int k =1;
		UINT l = 0;
		UINT m =0;
		UINT filelistoffest;
		offest.QuadPart = 64*512;
		status = ZwReadFile(Disk->CurrentDiskHandle,NULL,NULL,NULL,&io_stk,(PVOID)buffer,2,&offest,NULL);
		RtlCopyMemory(a.hexarray,buffer,2);
		KeWaitForSingleObject(&Disk->Volume->FileListLock,Executive,KernelMode,FALSE,NULL);
		for (UINT j = 0;j<a.shortnum;i++){
			buf = (BufStruct*)ExAllocatePoolWithTag(PagedPool,sizeof(BufStruct),'file');
			buf->buf = ExAllocatePoolWithTag(PagedPool,510,'lstb');
			buf->Length = 510;
			InsertTailList(&buf_list,&(buf->list_entry));
		}
		//读取记录
		offest.QuadPart = (64*512)+(a.shortnum)*512;
		for (LIST_ENTRY *pBuf = buf_list.Flink;pBuf != &buf_list;pBuf=pBuf->Flink){
			buf = CONTAINING_RECORD(pBuf,BufStruct,list_entry);
			status = ZwReadFile(Disk->CurrentDiskHandle,NULL,NULL,NULL,&io_stk,(buf->buf),510,&offest,NULL);
			offest.QuadPart += (510+2);
		}
		//分解记录
		temps[1] = ExAllocatePool(PagedPool,(a.shortnum)*512);
		//首先合并记录
		for (LIST_ENTRY *pBuf = buf_list.Flink;pBuf != &buf_list;pBuf=pBuf->Flink){
			buf = CONTAINING_RECORD(pBuf,BufStruct,list_entry);
			memcat(temps[1],buf->buf,l,buf->Length);
			ExFreePool(buf->buf);
			ExFreePool(buf);
			l += 510;
		}
		temps[4] = temps[1];
		//正式分解记录
		struct FItem
		{
			uchar size;
			uchar flag;
			uchar4 data_start;
			uchar4 data_size;
			time_t data_time;
			char name[255];
		};
		for (;k<l;)
		{
			//实现单链表式存储
			FItem Item;
			char size = *((char*)temps[1]);
			RtlCopyMemory(&Item,temps[1],size);
			size_t namelen = utf8_to_wchar(Item.name,Item.size-1-4-4-4,NULL,0,UTF8_IGNORE_ERROR);
			wchar_t *namebuffer = (wchar_t*)ExAllocatePool(PagedPool,namelen+1*sizeof(wchar_t));
			utf8_to_wchar(Item.name,Item.size-1-4-4-4,namebuffer,namelen+1*sizeof(wchar_t),UTF8_IGNORE_ERROR);
			namebuffer[namelen] = NULL;
			FileRecord *record = (FileRecord*)ExAllocatePool(PagedPool,sizeof(FileRecord));
			RtlZeroMemory(record,sizeof(FileRecord));
			RtlInitUnicodeString(&record->FileName[1],namebuffer);
			record->FileLength = Item.data_size;
			record->ModTime.StdModTime = Item.data_time;
			StdTimeToKeTime(record->ModTime.StdModTime,record->ModTime.KeModTime);
			record->offest.QuadPart = Item.data_start;
			record->OnDevice = Disk->Volume;
			record->OnVolume = Disk;
			InsertTailList(&Disk->Volume->FileListHead,&record->list_item);
		}
		//最后，释放内存
		ExFreePool(temps[4]);
		MakeTree(*Disk->Volume);
	}__except(EXCEPTION_EXECUTE_HANDLER){
		KeReleaseMutex(&Disk->Volume->FileListLock,FALSE);
	}
	return;
}

unsigned int NumChar(wchar_t *str,wchar_t *substr){
	int i;
	wchar_t *cp;
	
    	i = 0;
    	while(1) {
		cp = wcsstr(str,substr);
		if (cp != NULL){
        		str = ++cp;
			i++;
		}
		else
			break;
    	}
    	return i;
}

NTSTATUS SearchFileTree(UNICODE_STRING &FileName,FileRecord * Results[],int resultnum,FBFSVolume &Volume,FileRecord *Start){
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	static unsigned int count;
	if ((!Results) || (!resultnum))
	{
		return STATUS_INVALID_ADDRESS;
	}
	__try{
		if (Start=NULL){
			Start = Volume.RootDir;
		}
		for (FileRecord *Cur = Start;(Cur != Volume.RootDir) && (count<resultnum);Cur = Cur->DirStruct.next_sibling_file)
		{
			if (RtlEqualUnicodeString(&FileName,&(Cur->FileName[0]),FALSE))
			{
				Results[count] = Cur;
			}
			count++;
			if (Cur->DirStruct.IsDir)
			{
				status = SearchFileTree(FileName,Results+count,resultnum-count,Volume,Cur);
				if (!NT_SUCCESS(status))
				{
					*except = 0;//引发异常
				}
			}
		}
		if (count==0)
		{
			return STATUS_OBJECT_NAME_NOT_FOUND;
		}else{
			status = STATUS_SUCCESS;
		}
	}__except(EXCEPTION_EXECUTE_HANDLER){
		;
	}
	return status;
}

NTSTATUS MakeTree(FBFSVolume &Volume,FileRecord* Start,FileRecord* Root){
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	PAGED_CODE();
	__try{
		for (PLIST_ENTRY p = Volume.FileListHead.Flink;p != &Volume.FileListHead;p=p->Flink)
		{
			FileRecord * record;
			if (Start==NULL)
				record = CONTAINING_RECORD(p,FileRecord,list_item);
			else
				record = Start;
			wchar_t *cp[2] = {record->FileName->Buffer,record->FileName->Buffer};
			if (wcsstr(record->FileName[1].Buffer,L"/")!=NULL)
			{
				while (cp[0] == NULL)
				{
					cp[0] = wcsstr(cp[0],L"/");
					if (cp[0] != NULL)
					{
						cp[1] = cp[0];
					}
				}
				cp[1]++;
				RtlInitUnicodeString(&record->FileName[0],cp[1]);
				cp[0] = record->FileName->Buffer;
				FileRecord* cur[2] = {Root,Root};
				if (Start==NULL)
				{
					FileRecord* cur[2] = {Volume.RootDir,Volume.RootDir};
				}
				while (cp[0] == NULL)
				{
					cp[1] = wcsstr(cp[0],L"/");
					cp[0][(cp[1]-cp[0])] = NULL;
					//求出树链表最尾端的节点
					while (cur[0] != NULL)
					{
						cur[0] = cur[0]->DirStruct.next_sibling_file;
						if (cur[0] != NULL)
						{
							cur[1] = cur[0];
						}
					}
					UNICODE_STRING DirName = RTL_CONSTANT_STRING(cp[1]);
					FileRecord *Temp[2] = {Volume.RootDir,Volume.RootDir};
					bool found = false;
					while (Temp[0] != NULL)
					{
						if (RtlEqualUnicodeString(&DirName,&Temp[0]->FileName[0],TRUE)==TRUE)
						{
							found = true;
							break;
						}
						Temp[0] = Temp[0]->DirStruct.next_sibling_file;
					}
					if (found)
					{
						MakeTree(Volume,record,Temp[0]->DirStruct.lower_file);//递归
					}else{
						Temp[0]->DirStruct.next_sibling_file = (FileRecord *)ExAllocatePool(PagedPool,sizeof(FileRecord));
						Temp[0]->DirStruct.IsDir = true;
						Temp[0]->DirStruct.FatherDir = Temp[0];
						RtlInitEmptyUnicodeString(&Temp[0]->FileName[0],(PWCHAR)ExAllocatePool(NonPagedPool,DirName.MaximumLength),DirName.MaximumLength);
						MakeTree(Volume,record,Temp[0]->DirStruct.lower_file);//递归
					}
				}
			}else{
				RtlInitUnicodeString(&record->FileName[0],record->FileName->Buffer);
				FileRecord* cur[2] = {Volume.RootDir,Volume.RootDir};
				//求出树链表最尾端的节点
				while (cur[0] != NULL)
				{
					cur[0] = cur[0]->DirStruct.next_sibling_file;
					if (cur[0] != NULL)
					{
						cur[1] = cur[0];
					}
				}
				cur[1]->DirStruct.next_sibling_file = record;
				record->DirStruct.FatherDir = Volume.RootDir;
				record->DirStruct.IsDir = false;
			}
		}
		status = STATUS_SUCCESS;
	}
	__except(EXCEPTION_EXECUTE_HANDLER){
		;
	}
	return status;
}
